package com.dhl.aichat.controller;

import com.dhl.aichat.result.Result;
import lombok.RequiredArgsConstructor;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.client.RestTemplate;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/ollama")
@RequiredArgsConstructor
public class OllmaController {

    // Ollama默认API地址
    private static final String OLLAMA_API_URL = "http://localhost:11434/api/tags";

    private final RestTemplate restTemplate;

    // 获取本地ollma模型列表
    @GetMapping("/list")
    public Result<List<String>> list() {
        try {
            // 调用Ollama API
            ResponseEntity<Map> response = restTemplate.getForEntity(OLLAMA_API_URL, Map.class);

            if (response.getStatusCode() == HttpStatus.OK && response.getBody() != null) {
                Map<String, Object> responseBody = response.getBody();
                List<Map<String, Object>> models = (List<Map<String, Object>>) responseBody.get("models");

                List<String> modelNames = new ArrayList<>();
                if (models != null) {
                    for (Map<String, Object> model : models) {
                        modelNames.add((String) model.get("name"));
                    }
                }

                return Result.success(modelNames);
            } else {
                return Result.error("获取模型列表失败: " + response.getStatusCode());
            }
        } catch (Exception e) {
            return Result.error("获取模型列表异常: " + e.getMessage());
        }
    }
}
